import os
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torch.utils import data as Data

from handle_souce_data import MakeDataSet


# 获取数据集
class GetDataLoader:
    # 图片的均值
    MEAN = [0.162, 0.151, 0.139]
    # 图片的方差
    STD = [0.058, 0.052, 0.048]
    # 图片resize的大小
    IMG_SIZE = (224,224)
    # 图片的通道
    IMG_CHANNELS = 3
    # 单例模式：对象的地址信息
    INSTANCE = None

    def __new__(cls, *args, **kwargs):
        if cls.INSTANCE:
            return cls.INSTANCE
        else:

            cls.INSTANCE = super().__new__(cls)
            return cls.INSTANCE

    def __init__(self):
        # 图片的种类
        # 训练图片保存的目录
        TRAIN_PATH = os.path.join(MakeDataSet.save_path, MakeDataSet.train_dir_name)
        self.IMG_CLASSES: list = os.listdir(TRAIN_PATH)
        # 训练时：训练集和验证集
        self.train_and_val_dataloader: tuple = self._make_train_and_val_dataloader()
        # 验证时：测试数据集
        self.test_dataloader = self._make_test_dataloader()

    # 获取训练数据
    def _make_train_and_val_dataloader(self):
        # 均值/方差是算出来的，进行正态发布,mean和std有3个数值，说明图片是3通道的
        # 获取均值和均值方差
        # 先执行获取均值和均值方差 = img_mean_and_std()
        # 图片处理的大小
        if hasattr(self,'MEAN') and hasattr(self,'STD'):
            # 数据归一化处理
            normalize = transforms.Normalize(mean=self.MEAN, std=self.STD)
            # 图片的处理流程
            transform = transforms.Compose([transforms.Resize(self.IMG_SIZE), transforms.ToTensor(), normalize])
        else:
            # 图片的处理流程
            transform = transforms.Compose([transforms.Resize(self.IMG_SIZE), transforms.ToTensor()])
        # 数据集处理类
        make_cls = MakeDataSet
        # 自定义训练集的路径
        train_root_path = os.path.join(make_cls.save_path, make_cls.train_dir_name)
        train_data = ImageFolder(train_root_path, transform=transform)
        # print('数据的分类', train_data.class_to_idx, train_data.classes)
        # 数据集分割
        train_data, val_data = Data.random_split(
            train_data,
            [round(0.8 * len(train_data)), round(0.2 * len(train_data))]
        )

        # 训练数据加载
        train_dataloader = Data.DataLoader(
            dataset=train_data,
            batch_size=64,
            shuffle=True,
        )

        # 验证数据集加载
        val_dataloader = Data.DataLoader(
            dataset=val_data,
            batch_size=64,
            shuffle=True
        )
        return train_dataloader, val_dataloader
    # 获取测试集
    def _make_test_dataloader(self):
        # 均值/方差是算出来的，进行正态发布,mean和std有3个数值，说明图片是3通道的
        # 获取均值和均值方差
        # 先执行获取均值和均值方差 = img_mean_and_std()
        if hasattr(self, 'MEAN') and hasattr(self, 'STD'):
            # 数据归一化处理
            normalize = transforms.Normalize(mean=self.MEAN, std=self.STD)
            # 图片的处理流程
            transform = transforms.Compose([transforms.Resize(self.IMG_SIZE), transforms.ToTensor(), normalize])
        else:
            # 图片的处理流程
            transform = transforms.Compose([transforms.Resize(self.IMG_SIZE), transforms.ToTensor()])
        # 数据集处理类
        make_cls = MakeDataSet
        # 自定义训练集的路径
        test_root_path = os.path.join(make_cls.save_path, make_cls.test_dir_name)
        test_data = ImageFolder(test_root_path, transform=transform)
        # print('数据的分类', test_data.class_to_idx, test_data.classes)

        test_dataloader = DataLoader(
            dataset=test_data, batch_size=1, shuffle=True, num_workers=0
        )

        return test_dataloader

if __name__ == '__main__':
    # 数据集的情况
    loader = GetDataLoader()
    # 数据集中图片的通道
    channels = loader.IMG_CHANNELS
    # 图片要resize的大小
    img_size = loader.IMG_SIZE
    # 图片的类型，[分类名1,分类名2]
    img_classes = loader.IMG_CLASSES
    # 训练数据加载和验证数据加载
    train_dataloader, val_dataloader = loader.train_and_val_dataloader
    # 测试数据加载
    test_dataloader = loader.test_dataloader